#include "lcm_publish_onnx.h"
#include <iostream>
#include <valarray>


using namespace std;

RL_Xdog xdog_rl;

 
int RL_Xdog::init_onnx(){

    // 指定模型文件路径
    policy_path = "/home/saw/C++/sim2sim/onnx/bset.onnx";
    
    // 1. 创建ONNX Runtime环境对象，设置日志级别为WARNING
    Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "ONNX");
    
    // 2. 创建会话选项对象，用于配置推理参数
    Ort::SessionOptions session_options;
    
    // 3. 创建默认内存分配器，用于管理张量内存
    Ort::AllocatorWithDefaultOptions allocator;
    
    // 4. 加载ONNX模型，创建会话对象
    onnx_session = new Ort::Session(env, policy_path.c_str(), session_options);
    
    // 5. 获取模型输入节点信息
    input_nodes_num = onnx_session->GetInputCount();  // 输入节点数量
    std::cout << "Number of input nodes: " << input_nodes_num << std::endl; //1个节点数量
    
    for (int i = 0; i < input_nodes_num; ++i) {
        // 获取输入节点名称和形状
        auto input_name = onnx_session->GetInputNameAllocated(i, allocator);
        input_node_names.push_back(input_name.get());  // 存储输入节点名称
        auto inputShapeInfo = onnx_session->GetInputTypeInfo(i)
                                .GetTensorTypeAndShapeInfo()
                                .GetShape();  // 获取输入形状

        std::cout << "Input shape: " << inputShapeInfo.size() << " dimensions" << std::endl;

        // 打印输入形状的具体维度（如batch、通道等）
        for (size_t j = 0; j < inputShapeInfo.size(); ++j) {
            std::cout << "inputShapeInfo["<< j << "]: " << inputShapeInfo[j] << std::endl;
        }
    }
    
    // 6. 获取模型输出节点信息（逻辑同输入节点）
    output_nodes_num = onnx_session->GetOutputCount();
    std::cout << "Number of output nodes: " << output_nodes_num << std::endl;
    for (int i = 0; i < output_nodes_num; ++i) {
        auto output_name = onnx_session->GetOutputNameAllocated(i, allocator);
        output_node_names.push_back(output_name.get());
        auto outShapeInfo = onnx_session->GetOutputTypeInfo(i)
                             .GetTensorTypeAndShapeInfo()
                             .GetShape();
        std::cout << "Output shape: " << outShapeInfo.size() << " dimensions" << std::endl;
        for (size_t j = 0; j < outShapeInfo.size(); ++j) {
            std::cout << "outShapeInfo["<< j << "]: " << outShapeInfo[j] << std::endl;
        }
    }
    
    // 7. 确定输入/输出节点的名称（假设模型只有1个输入和1个输出）
    input_name = {input_node_names[0].c_str()};
    output_name = {output_node_names[0].c_str()};
    
    // 8. 初始化动作向量（用于存储推理结果）
    action = std::vector<double>(12, 0.0);  // 假设输出动作维度为12，所有元素默认值为0.0f
    obs = std::vector<float>(45, 0.0f);    // 假设输入观测数据维度为45，所有元素默认值为0.0f
}



void RL_Xdog::handleMessage(const lcm::ReceiveBuffer *rbuf, const std::string &chan,
                       const my_lcm::Request *request)//获取机器人反馈
{              

    // 2. 模拟观测数据（45维向量，归一化到[0,1]）
    std::vector<float> obs_test(45, 0.0f);

    // 将接收到的数据压入到obs序列
    obs_test[0]=request->omega[0]*omega_scale;
    obs_test[1]=request->omega[1]*omega_scale;
    obs_test[2]=request->omega[2]*omega_scale;

    obs_test[3]=request->gravity_orientation[0];
    obs_test[4]=request->gravity_orientation[1];
    obs_test[5]=request->gravity_orientation[2];

    obs_test[6]=cmd_x*cmd_lin_vel;//控制指令x
    obs_test[7]=cmd_y*cmd_lin_vel;//控制指令y
    obs_test[8]=cmd_z*cmd_ang_vel;//控制指令z */

    //关节角度
    for (int i = 9; i < 9 + 12; ++i){
        obs_test[i]= (request->q[i-9]  - init_pos[i-9])* dof_pos_scale;
    }

    //关节角速度
    for (int i = 9 + 12; i < 9+12+12; ++i){
        obs_test[i]= request->dq[i-9-12] * dof_vel_scale;
    }

    //action
    for (int i = 9+12+12; i < 45; ++i){
        obs_test[i] = action[i-9-12-12]; 
    } 

    obs = obs_test;  // 赋值给成员变量

    // 1. 创建输入张量：根据输入形状和数据生成内存中的张量
    auto memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);

    // 输入形状为[1, 45]（batch=1，数据长度45），数据来自obs成员变量
    std::array<int64_t, 1> input_shape{static_cast<int64_t>(obs.size())};

    Ort::Value input_tensor = Ort::Value::CreateTensor<float>(
        memory_info, obs.data(), obs.size(), input_shape.data(), input_shape.size());

    std::cout << "obs size :" << obs.size() << std::endl;

    // 2. 执行推理：传入输入张量和名称，获取输出张量
    std::vector<Ort::Value> output_tensor;  // 存储输出张量

    output_tensor = onnx_session->Run(
        Ort::RunOptions{nullptr},  // 运行选项（可设置超时等参数）
        input_name.data(),         // 输入节点名称数组
        &input_tensor,             // 输入张量指针
        1,                         // 输入数量
        output_name.data(),        // 输出节点名称数组
        output_name.size());       // 输出数量
    
    // 3. 提取输出张量的数据并转换为double类型
    float* float_array = output_tensor[0].GetTensorMutableData<float>();
    for (size_t i = 0; i < action.size(); ++i) {
        action[i] = static_cast<double>(float_array[i]);  // 转换并存储结果
        printf("action %d  :  %.2f \n",i,action[i]);
    }

}


int main(int argc, char** argv) {

   lcm::LCM lcm;

    if (!lcm.good()) {
        std::cerr << "LCM initialization failed" << std::endl;
        return 1;
    }
    
    //加载策略
    xdog_rl.init_onnx(); 

    //创建订阅者
    lcm.subscribe("LCM_OBS", &RL_Xdog::handleMessage, &xdog_rl);

    while (1){
        my_lcm::Response msg;
        for(int i=0;i<12;i++){
            msg.q_exp[i]=xdog_rl.action[i];
        }
            
        lcm.publish("LCM_ACTION", &msg);
        //std::cout << "Message robot state published!" <<std::endl;
        lcm.handle();
    } 

    return 0;
}

 


